from pydantic import BaseModel
from dataclasses import dataclass

@dataclass
class PythonFunction:
    function_defined: str = ""
    function_body: str = ""
    function_args: str | None = ""

@dataclass
class AirflowOperator:
    operator_name: str = None
    operator_type: str = None
    python_function_defined: str | None = None
    python_function_args: list[dict] | None = None
    spark_application: str | None = None

@dataclass
class UserDataSource:
    datasource_type: str = None
    datasource_convert: str | None = None
    bucket_name: str | None = ""
    object_name: str | None = ""
    convert_result_name: str | None = ""

@dataclass
class BuildIDAGItem(BaseModel):
    user_id: str = "yz"
    operators: list[AirflowOperator] | None = None
    dag_sequence: list[str] | None = None
    dag_sequence_tree: dict | None = None
    python_functions: list[PythonFunction] | None = None
    spark_applications: list[str] | None = None
    python_dependencies: list[str] | None = None
    user_data_source: UserDataSource | None = None

@dataclass
class TriggerDAGItem(BaseModel):
    dag_id: str = ""

@dataclass
class ActivateDAGItem(BaseModel):
    dag_id: str = ""